import matplotlib

matplotlib.use('Agg')  # 使用非交互式后端，适合生成图片到文件
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd
from io import BytesIO, StringIO
import sys
import traceback
import time
from typing import Dict, Any, Optional


class SafeChartGenerator:
    """
    安全的图表生成器

    在受限环境中执行用户代码生成图表，防止恶意代码执行
    支持 matplotlib、seaborn、numpy 和 pandas 库
    """

    def __init__(self, default_timeout: int = 10):
        """
        初始化安全图表生成器

        Args:
            default_timeout: 默认代码执行超时时间（秒）
        """
        self.default_timeout = default_timeout

        # 定义允许使用的模块
        self.allowed_modules = {
            'matplotlib.pyplot': plt,
            'seaborn': sns,
            'matplotlib': matplotlib,
            'numpy': np,
            'pandas': pd
        }

        # 创建受限的命名空间
        self.namespace = {
            # 导入允许的模块
            'plt': plt,
            'sns': sns,
            'matplotlib': matplotlib,
            'np': np,
            'pd': pd,
            # 允许一些基本函数
            'print': print
        }

    def reset_environment(self) -> None:
        """重置matplotlib环境状态，确保每次执行都是干净的环境"""
        plt.close('all')  # 关闭所有之前的图形
        # 清理可能的状态
        plt.clf()
        plt.cla()

    def _safe_import(self, module_name: str):
        """
        安全导入模块，只允许白名单中的模块

        Args:
            module_name: 模块名称

        Returns:
            导入的模块对象

        Raises:
            ImportError: 当尝试导入不允许的模块时
        """
        if module_name in self.allowed_modules:
            return self.allowed_modules[module_name]
        raise ImportError(f"不允许导入模块: {module_name}")

    def execute_code(self, code: str, timeout: Optional[int] = None) -> Dict[str, Any]:
        """
        在安全环境中执行代码并返回生成的图片

        Args:
            code: 要执行的Python代码字符串
            timeout: 执行超时时间(秒)，如果为None则使用默认值

        Returns:
            包含执行结果的字典，格式为:
            {
                'success': bool,  # 是否成功
                'image': BytesIO,  # 生成的图片(成功时)
                'output': str,  # 执行输出
                'error': str,  # 错误信息(失败时)
                'traceback': str  # 详细错误追踪(失败时)
            }
        """
        if timeout is None:
            timeout = self.default_timeout

        # 重置环境
        self.reset_environment()

        # 初始化变量
        old_stdout = sys.stdout
        redirected_output = StringIO()

        try:
            # 重定向标准输出，捕获print内容
            sys.stdout = redirected_output

            # 记录开始时间，用于超时检查
            start_time = time.time()

            # 定义超时检查函数
            def check_timeout():
                if time.time() - start_time > timeout:
                    raise TimeoutError(f"代码执行超时，超过{timeout}秒")

            # 在受限命名空间中执行代码
            compiled_code = compile(code, '<user_code>', 'exec')
            exec(compiled_code, self.namespace)
            check_timeout()

            # 检查是否有生成的图形
            if plt.get_fignums():
                # 将图形保存到字节流
                img_buffer = BytesIO()

                # 如果有多个图形，只保存最后一个活动图形
                # 或者保存当前活动的图形
                plt.savefig(img_buffer, format='png', bbox_inches='tight', dpi=300)
                img_buffer.seek(0)
                result_image = img_buffer
            else:
                # 没有生成图形
                raise RuntimeError("代码没有生成任何图形")

            # 恢复标准输出
            sys.stdout = old_stdout
            output = redirected_output.getvalue()

            return {
                'success': True,
                'image': result_image,
                'output': output
            }

        except Exception as e:
            # 恢复标准输出
            sys.stdout = old_stdout
            output = redirected_output.getvalue()

            # 获取详细的错误信息
            error_msg = traceback.format_exc()

            return {
                'success': False,
                'error': str(e),
                'traceback': error_msg,
                'output': output
            }
        finally:
            # 确保标准输出恢复
            sys.stdout = old_stdout
            # 重置环境，为下一次执行做准备
            self.reset_environment()

    def generate_chart_to_file(self, code: str, output_path: str, timeout: Optional[int] = None) -> Dict[str, Any]:
        """
        执行代码并将生成的图片保存到文件

        Args:
            code: 要执行的Python代码字符串
            output_path: 输出图片文件路径
            timeout: 执行超时时间(秒)

        Returns:
            包含执行结果的字典
        """
        result = self.execute_code(code, timeout)

        if result['success']:
            try:
                with open(output_path, 'wb') as f:
                    f.write(result['image'].read())
                result['output_path'] = output_path
            except Exception as e:
                result['success'] = False
                result['error'] = f"保存文件失败: {str(e)}"

        return result


def safe_execute_image_code(code: str, timeout: int = 10) -> Dict[str, Any]:
    """
    安全执行生成图片的Python代码的便捷函数

    Args:
        code: 要执行的Python代码
        timeout: 超时时间(秒)

    Returns:
        包含执行结果的字典
    """
    generator = SafeChartGenerator(default_timeout=timeout)
    return generator.execute_code(code, timeout)
